from ultralytics import YOLO
import os
from datetime import datetime
import cv2
import numpy as np
from openvino.runtime import Core
import time
import signal
import sys

# 保存检测图片开关
save_pic = True
# 加载模型
model = YOLO('models/yolov8n-seg.pt')

# OpenVINO 初始化
core = Core()
road_model_path = r"models\road-segmentation-adas-0001\FP32\road-segmentation-adas-0001.xml"
road_model = core.read_model(model=road_model_path)
road_compiled = core.compile_model(model=road_model, device_name="CPU")
road_input_layer = road_compiled.input(0)
road_output_layer = road_compiled.output(0)

# 获取道路分割模型的输入尺寸
_, _, h_seg, w_seg = road_input_layer.shape

# RTSP视频流配置
RTSP_URL = "rtsp://admin:wkkj1234@192.168.200.74:554/Steaming/Channels/1"  # 替换为你的RTSP流地址
VEHICLE_CLASSES = [2, 5, 7]  # 只检测 car=2, bus=5, truck=7
CLASS_NAMES = {2: 'car', 5: 'bus', 7: 'truck'}

# 全局退出标志
should_exit = False

# 压线检测相关参数
VIOLATION_THRESHOLD = 0.02  # 压线面积阈值，超过此比例认为压线
BOTTOM_REGION_HEIGHT_RATIO = 0.3  # 车辆底部区域高度占总高度的比例（30%）
BOTTOM_REGION_WIDTH_RATIO = 0.9  # 车辆底部区域宽度占总宽度的比例（80%）
last_violation_time = {}  # 记录每个车辆上次压线时间，用于防抖
VIOLATION_COOLDOWN = 3  # 压线告警冷却时间（秒）

# 调试参数
SHOW_DETECTION_BOX = True  # 是否显示检测区域矩形框


# ROI区域配置 - 自定义梯形4个点的坐标
def create_trapezoid_roi_mask(frame_shape, top_left, top_right, bottom_right, bottom_left):
    """
    创建自定义梯形ROI检测区域掩码
    Args:
        frame_shape: 帧的形状 (height, width, channels)
        top_left: 梯形左上角点 (x, y)
        top_right: 梯形右上角点 (x, y)
        bottom_right: 梯形右下角点 (x, y)
        bottom_left: 梯形左下角点 (x, y)
    Returns:
        roi_mask: 梯形ROI掩码
    """
    h, w = frame_shape[:2]

    # 验证坐标是否在图像范围内
    def validate_point(point, width, height):
        x, y = point
        return max(0, min(width - 1, x)), max(0, min(height - 1, y))

    top_left = validate_point(top_left, w, h)
    top_right = validate_point(top_right, w, h)
    bottom_right = validate_point(bottom_right, w, h)
    bottom_left = validate_point(bottom_left, w, h)

    # 创建梯形ROI掩码
    roi_mask = np.zeros((h, w), dtype=np.uint8)
    points = np.array([top_left, top_right, bottom_right, bottom_left], np.int32)
    cv2.fillPoly(roi_mask, [points], 1)

    return roi_mask


# 全局ROI掩码
roi_mask = None


def signal_handler(sig, frame):
    """处理Ctrl+C信号"""
    global should_exit
    print('\n接收到退出信号，正在关闭...')
    should_exit = True


def check_window_closed():
    """检查窗口是否被关闭"""
    try:
        if cv2.getWindowProperty('Vehicle + Road Markings Detection', cv2.WND_PROP_VISIBLE) < 1:
            return True
        return False
    except:
        # 如果窗口不存在，返回False
        return False


def get_road_segmentation_mask(frame):
    """
    使用道路分割模型获取道路标记掩码
    """
    # 预处理图像
    processed_frame = preprocess_image_for_segmentation(frame)

    # 调整尺寸并进行推理
    resized_seg = cv2.resize(processed_frame, (w_seg, h_seg))
    input_seg = np.expand_dims(resized_seg.transpose(2, 0, 1), axis=0).astype(np.float32)

    # 推理
    infer_result = road_compiled([input_seg])
    tensor_output = infer_result[road_output_layer]
    seg_probs = np.array(tensor_output)
    seg_class = np.argmax(seg_probs[0], axis=0).astype(np.uint8)

    # 缩放回原图尺寸
    orig_h, orig_w = frame.shape[:2]
    seg_class_resized = cv2.resize(seg_class, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)

    # 创建道路标记掩码（类别3）
    lane_marking_mask = (seg_class_resized == 3).astype(np.uint8)

    # 应用ROI掩码，只保留ROI区域的道路标记
    global roi_mask
    if roi_mask is not None:
        lane_marking_mask = cv2.bitwise_and(lane_marking_mask, roi_mask)

    return lane_marking_mask


def preprocess_image_for_segmentation(image):
    """
    对输入图像进行预处理，增强对比度以改善分割效果
    """
    processed = image.copy()

    # 自适应方法：根据图像亮度选择最佳预处理
    gray = cv2.cvtColor(processed, cv2.COLOR_BGR2GRAY)
    brightness = np.mean(gray)

    if brightness < 80:  # 暗图像
        # Gamma校正提亮
        gamma = 1.8
        inv_gamma = 1.0 / gamma
        table = np.array([((i / 255.0) ** inv_gamma) * 255 for i in np.arange(0, 256)]).astype("uint8")
        processed = cv2.LUT(processed, table)

        # CLAHE增强对比度
        lab = cv2.cvtColor(processed, cv2.COLOR_BGR2LAB)
        lab_planes = list(cv2.split(lab))
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
        lab_planes[0] = clahe.apply(lab_planes[0])
        lab = cv2.merge(lab_planes)
        processed = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)

    elif brightness > 200:  # 过亮图像
        # 过亮图像主要增强对比度
        lab = cv2.cvtColor(processed, cv2.COLOR_BGR2LAB)
        lab_planes = list(cv2.split(lab))
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        lab_planes[0] = clahe.apply(lab_planes[0])
        lab = cv2.merge(lab_planes)
        processed = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)

    else:  # 正常亮度
        # 正常图像使用标准CLAHE
        lab = cv2.cvtColor(processed, cv2.COLOR_BGR2LAB)
        lab_planes = list(cv2.split(lab))
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        lab_planes[0] = clahe.apply(lab_planes[0])
        lab = cv2.merge(lab_planes)
        processed = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)

    return processed


def create_vehicle_bottom_detection_region(bbox, frame_shape, height_ratio=BOTTOM_REGION_HEIGHT_RATIO,
                                           width_ratio=BOTTOM_REGION_WIDTH_RATIO):
    """
    在车辆底部创建一个矩形检测区域
    Args:
        bbox: 车辆边界框 (x1, y1, x2, y2)
        frame_shape: 原始帧的形状 (height, width, channels)
        height_ratio: 底部区域高度占车辆总高度的比例
        width_ratio: 底部区域宽度占车辆总宽度的比例
    Returns:
        bottom_region_mask: 底部检测区域掩码
        detection_bbox: 检测区域边界框 (x1, y1, x2, y2)
    """
    x1, y1, x2, y2 = bbox
    h, w = frame_shape[:2]

    # 计算车辆中心点和尺寸
    vehicle_center_x = (x1 + x2) // 2
    vehicle_center_y = (y1 + y2) // 2
    vehicle_width = x2 - x1
    vehicle_height = y2 - y1

    # 计算底部区域的尺寸
    bottom_width = int(vehicle_width * width_ratio)
    bottom_height = int(vehicle_height * height_ratio)

    # 计算底部区域的中心点（在车辆底部）
    bottom_center_y = y2 - bottom_height // 2

    # 计算底部区域的坐标
    bottom_x1 = max(0, vehicle_center_x - bottom_width // 2)
    bottom_y1 = max(0, bottom_center_y - bottom_height // 2)
    bottom_x2 = min(w, vehicle_center_x + bottom_width // 2)
    bottom_y2 = min(h, bottom_center_y + bottom_height // 2) + 100  # 向后扩展检测区域，如果车后方很近的地方有实线，认为其压线

    # 创建底部区域掩码
    bottom_region_mask = np.zeros((h, w), dtype=np.uint8)
    bottom_region_mask[bottom_y1:bottom_y2, bottom_x1:bottom_x2] = 1

    # 应用ROI掩码，只保留ROI区域内的检测区域
    global roi_mask
    if roi_mask is not None:
        bottom_region_mask = cv2.bitwise_and(bottom_region_mask, roi_mask)

    # 返回检测区域掩码和检测区域边界框
    detection_bbox = (bottom_x1, bottom_y1, bottom_x2, bottom_y2)

    return bottom_region_mask, detection_bbox

def check_vehicle_violation(bbox, lane_marking_mask, frame_shape):
    """
    检查车辆底部检测区域是否压线
    """
    if lane_marking_mask is None:
        return False, bbox

    # 创建车辆底部检测区域
    bottom_region_mask, detection_bbox = create_vehicle_bottom_detection_region(bbox, frame_shape)

    # 确保两个掩码尺寸相同
    if bottom_region_mask.shape != lane_marking_mask.shape:
        bottom_region_mask = cv2.resize(bottom_region_mask.astype(np.uint8),
                                        (lane_marking_mask.shape[1], lane_marking_mask.shape[0]))
        bottom_region_mask = bottom_region_mask.astype(np.uint8)

    # 计算底部检测区域与道路标记的交集
    intersection = cv2.bitwise_and(bottom_region_mask, lane_marking_mask)

    # 计算交集面积和底部检测区域面积
    intersection_area = np.sum(intersection)
    bottom_region_area = np.sum(bottom_region_mask)

    if bottom_region_area == 0:
        return False, detection_bbox

    # 计算压线比例
    overlap_ratio = intersection_area / bottom_region_area

    # 添加详细调试信息
    print(f"检测区域坐标: {detection_bbox}")
    print(f"底部检测区域面积: {bottom_region_area}")
    print(f"交集面积: {intersection_area}")
    print(f"压线比例: {overlap_ratio}")
    print(f"车辆原始bbox: {bbox}")

    # 判断是否压线
    is_violation = overlap_ratio > VIOLATION_THRESHOLD
    if is_violation:
        print(f"⚠️  检测到压线! 比例: {overlap_ratio:.4f} > {VIOLATION_THRESHOLD}")
        print(f"   车辆bbox: {bbox}")
        print(f"   检测区域: {detection_bbox}")
        print(
            f"   扩展了多少像素: {detection_bbox[3] - (bbox[3] + int((bbox[3] - bbox[1]) * BOTTOM_REGION_HEIGHT_RATIO))}")

    return is_violation, detection_bbox


def combine_vehicle_and_road_masks(original_frame, results, lane_marking_mask):
    """
    将车辆掩码和道路标记掩码合并到同一张图片上
    """
    # 确保lane_marking_mask与原图尺寸一致
    if lane_marking_mask.shape[:2] != original_frame.shape[:2]:
        lane_marking_mask = cv2.resize(lane_marking_mask, (original_frame.shape[1], original_frame.shape[0]))

    # 创建结果图像副本
    result_frame = original_frame.copy()

    # 创建车辆掩码的透明叠加层
    vehicle_mask_layer = np.zeros_like(result_frame)

    vehicle_count = 0
    violations = []  # 存储压线信息

    # 遍历检测结果，先进行压线检测并存储结果
    vehicle_violations = {}  # 存储每辆车的压线检测结果

    if results.boxes is not None:
        for i, box in enumerate(results.boxes):
            cls_id = int(box.cls[0])
            conf = box.conf[0].item()

            # 获取车辆边界框
            x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())

            # 检查车辆是否在ROI区域内
            vehicle_center_x = (x1 + x2) // 2
            vehicle_center_y = (y1 + y2) // 2

            global roi_mask
            if roi_mask is not None:
                # 检查车辆中心点是否在ROI区域内
                if roi_mask[vehicle_center_y, vehicle_center_x] == 0:
                    continue  # 车辆不在ROI区域内，跳过检测

            # 只处理车辆类别
            if cls_id in VEHICLE_CLASSES and conf > 0.5:
                vehicle_count += 1

                # 只进行一次压线检测
                is_violating, detection_bbox = check_vehicle_violation((x1, y1, x2, y2), lane_marking_mask,
                                                                       original_frame.shape)

                # 生成车辆唯一ID（基于边界框位置）
                vehicle_id = f"{cls_id}_{x1}_{y1}_{x2}_{y2}"

                # 检查是否在冷却时间内
                current_time = time.time()
                if vehicle_id in last_violation_time:
                    if current_time - last_violation_time[vehicle_id] < VIOLATION_COOLDOWN:
                        is_violating = False  # 在冷却时间内，不重复告警
                    else:
                        last_violation_time[vehicle_id] = current_time
                else:
                    last_violation_time[vehicle_id] = current_time

                # 存储压线检测结果
                vehicle_violations[i] = {
                    'is_violating': is_violating,
                    'detection_bbox': detection_bbox
                }

                # 存储压线信息
                if is_violating:
                    violations.append({
                        'vehicle_id': vehicle_id,
                        'class_id': cls_id,
                        'bbox': (x1, y1, x2, y2),
                        'class_name': CLASS_NAMES[cls_id],
                        'detection_bbox': detection_bbox
                    })

                # 为不同车辆分配不同颜色
                if is_violating:
                    color = (0, 0, 255)  # 压线车辆用红色
                else:
                    color = get_color_for_class(cls_id)

                # 获取车辆掩码（如果存在）
                if results.masks is not None and i < len(results.masks.data):
                    mask = results.masks.data[i].cpu().numpy()
                    # 调整掩码尺寸到原图大小
                    mask_resized = cv2.resize(mask, (original_frame.shape[1], original_frame.shape[0]))

                    # 创建当前车辆的掩码层
                    mask_bool = mask_resized > 0.3
                    current_vehicle_layer = np.zeros_like(vehicle_mask_layer)
                    current_vehicle_layer[mask_bool] = color

                    # 将当前车辆掩码添加到总层
                    vehicle_mask_layer = cv2.addWeighted(vehicle_mask_layer, 1,
                                                         current_vehicle_layer, 1, 0)

    # 将车辆掩码叠加到原图上
    alpha = 0.5  # 透明度
    result_frame = cv2.addWeighted(result_frame, 1, vehicle_mask_layer, alpha, 0)

    # 将道路标记掩码叠加到图像上（用红色表示）
    road_marking_layer = np.zeros_like(result_frame)
    road_marking_layer[lane_marking_mask == 1] = [0, 0, 255]  # 红色

    # 应用道路标记掩码
    result_frame = cv2.addWeighted(result_frame, 1, road_marking_layer, 0.3, 0)

    # 绘制ROI区域边界（可选）
    if SHOW_DETECTION_BOX and roi_mask is not None:
        # 找到ROI区域的轮廓
        contours, _ = cv2.findContours(roi_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for contour in contours:
            cv2.polylines(result_frame, [contour], True, (255, 255, 0), 2)  # 青色边界

    # 绘制边界框和标签（使用已存储的压线检测结果）
    if results.boxes is not None:
        for i, box in enumerate(results.boxes):
            cls_id = int(box.cls[0])
            conf = box.conf[0].item()

            # 获取车辆边界框
            x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())

            # 检查车辆是否在ROI区域内
            vehicle_center_x = (x1 + x2) // 2
            vehicle_center_y = (y1 + y2) // 2

            if roi_mask is not None:
                # 检查车辆中心点是否在ROI区域内
                if roi_mask[vehicle_center_y, vehicle_center_x] == 0:
                    continue  # 车辆不在ROI区域内，跳过检测

            if cls_id in VEHICLE_CLASSES and conf > 0.5:
                # 使用已存储的压线检测结果，不再重复调用check_vehicle_violation
                violation_info = vehicle_violations.get(i, {'is_violating': False, 'detection_bbox': (x1, y1, x2, y2)})
                is_violating = violation_info['is_violating']
                detection_bbox = violation_info['detection_bbox']

                if is_violating:
                    color = (0, 0, 255)  # 压线车辆边界框用红色
                    border_thickness = 4  # 压线车辆边界框更粗
                else:
                    color = get_color_for_class(cls_id)
                    border_thickness = 2

                # 绘制车辆边界框
                cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, border_thickness)

                # 如果启用调试模式，绘制检测区域矩形框
                if SHOW_DETECTION_BOX:
                    detection_color = (255, 255, 0)  # 青色表示检测区域
                    detection_thickness = 2
                    det_x1, det_y1, det_x2, det_y2 = detection_bbox
                    cv2.rectangle(result_frame, (det_x1, det_y1), (det_x2, det_y2), detection_color,
                                  detection_thickness)

                    # 在检测区域上方显示检测区域信息
                    cv2.putText(result_frame, f'Detect: {det_x2 - det_x1}x{det_y2 - det_y1}',
                                (det_x1, det_y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.4, detection_color, 1)

                # 绘制标签
                label = f"{CLASS_NAMES[cls_id]} {conf:.2f}"
                if is_violating:
                    label += " VIOLATION"

                label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
                cv2.rectangle(result_frame, (x1, y1 - label_size[1] - 10),
                              (x1 + label_size[0], y1), color, -1)
                cv2.putText(result_frame, label, (x1, y1 - 5),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)

    return result_frame, vehicle_count, violations


def create_video_capture(rtsp_url):
    """
    创建并配置视频捕获对象
    """
    print(f"尝试连接RTSP流: {rtsp_url}")

    # 创建视频捕获对象
    cap = cv2.VideoCapture()

    # 设置RTSP参数以提高稳定性
    cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)  # 减少缓冲区，降低延迟
    cap.set(cv2.CAP_PROP_FPS, 15)  # 设置期望的帧率
    cap.set(cv2.CAP_PROP_OPEN_TIMEOUT_MSEC, 10000)  # 设置打开超时
    cap.set(cv2.CAP_PROP_READ_TIMEOUT_MSEC, 10000)  # 设置读取超时

    # 尝试打开RTSP流
    success = cap.open(rtsp_url)

    if success:
        print("RTSP流连接成功!")
        return cap
    else:
        print("RTSP流连接失败!")
        return None


def process_rtsp_stream(rtsp_url, output_dir, process_interval=10):
    """
    处理RTSP视频流并进行车辆检测
    Args:
        rtsp_url: RTSP流地址
        output_dir: 输出目录
        process_interval: 处理间隔（每N帧处理一次）
    """
    global should_exit, roi_mask

    cap = create_video_capture(rtsp_url)
    if cap is None:
        return

    frame_count = 0
    processed_count = 0
    saved_count = 0
    consecutive_error_count = 0
    max_consecutive_errors = 5  # 最大连续错误次数

    print(f"成功连接RTSP流: {rtsp_url}")
    print("开始处理视频流...")

    # 创建窗口并设置大小为1066x600
    cv2.namedWindow('Vehicle + Road Markings Detection', cv2.WINDOW_NORMAL)
    cv2.resizeWindow('Vehicle + Road Markings Detection', 1066, 600)

    while not should_exit:
        try:
            ret, frame = cap.read()

            if not ret or frame is None or frame.size == 0:
                print(f"无法读取视频帧，可能流已断开 (连续错误: {consecutive_error_count})")
                consecutive_error_count += 1

                if consecutive_error_count >= max_consecutive_errors:
                    print("连续错误次数过多，尝试重新连接...")
                    cap.release()
                    time.sleep(2)  # 等待2秒后重连
                    cap = create_video_capture(rtsp_url)
                    if cap is None:
                        print("重新连接失败，退出...")
                        break
                    consecutive_error_count = 0
                    continue
                else:
                    time.sleep(0.1)  # 短暂等待
                    continue
            else:
                consecutive_error_count = 0  # 重置错误计数

            # 初始化ROI掩码（在第一帧时创建）
            if roi_mask is None:
                h, w = frame.shape[:2]
                # 默认梯形ROI坐标 (可根据需要调整)
                # top_left = (w // 4, h // 2)  # 左上角
                # top_right = (3 * w // 4, h // 2)  # 右上角
                top_left = (1300, 345)  # 左上角
                top_right = (1450, 345)  # 右上角
                # bottom_right = (w, h)  # 右下角
                # bottom_left = (0, h)  # 左下角
                bottom_right = (483, 732)  # 右下角
                bottom_left = (1881, 732)  # 左下角
                roi_mask = create_trapezoid_roi_mask(frame.shape, top_left, top_right, bottom_right, bottom_left)

            frame_count += 1

            # 每隔process_interval帧处理一次
            if frame_count % process_interval == 0:
                processed_count += 1
                print(f"\n处理第 {frame_count} 帧...")

                try:
                    # 获取道路标记掩码
                    lane_marking_mask = get_road_segmentation_mask(frame)

                    # 使用predict方法进行预测，确保掩码质量
                    results = model.predict(
                        source=frame,
                        imgsz=1280,  # 增大推理尺寸提高精度
                        conf=0.6,  # 置信度阈值
                        iou=0.5,  # IOU阈值
                        retina_masks=True,  # 启用高精度掩码
                        stream=False,  # 流模式优化
                        verbose=False,  # 关闭详细输出
                        save=False  # 不自动保存，我们自己处理
                    )

                    # 处理检测结果并合并道路标记掩码
                    visualized_frame, vehicle_count, violations = combine_vehicle_and_road_masks(
                        frame, results[0], lane_marking_mask)

                    # 检查压线情况并输出告警
                    if violations:
                        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                        print(f"[{current_time}] 检测到压线告警:")
                        for violation in violations:
                            print(
                                f"  - {violation['class_name']} 压线! 位置: ({violation['bbox'][0]}, {violation['bbox'][1]})")
                            print(f"    检测区域: {violation['detection_bbox']}")

                        # 如果有压线情况，保存图片
                        if save_pic:
                            save_segmented_image(visualized_frame, frame_count, output_dir, "violation")
                            saved_count += 1
                            print(f"压线告警图片已保存")
                    else:
                        print(f"未检测到压线")

                    # 只有检测到车辆时才保存图片
                    if vehicle_count > 0 and not violations:  # 没有压线但有车辆时也保存
                        if save_pic:
                            save_segmented_image(visualized_frame, frame_count, output_dir, "combined")
                            saved_count += 1
                            print(f"检测到 {vehicle_count} 辆车辆，已保存图片")
                    elif vehicle_count == 0:
                        print(f"未检测到车辆，跳过保存")

                    # 实时显示结果
                    display_results(visualized_frame, frame_count, vehicle_count, violations)

                except Exception as e:
                    print(f"处理第 {frame_count} 帧时出错: {e}")
                    import traceback
                    traceback.print_exc()
                    continue

            # 检查退出条件
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                print("检测到 'q' 键，退出...")
                should_exit = True
                break

            # 检查窗口是否被关闭
            if check_window_closed():
                print("检测到窗口关闭，退出...")
                should_exit = True
                break

        except KeyboardInterrupt:
            print("\n接收到键盘中断信号，正在退出...")
            should_exit = True
            break
        except Exception as e:
            print(f"读取帧时出错: {e}")
            import traceback
            traceback.print_exc()
            consecutive_error_count += 1
            if consecutive_error_count >= max_consecutive_errors:
                print("连续错误次数过多，尝试重新连接...")
                cap.release()
                time.sleep(2)
                cap = create_video_capture(rtsp_url)
                if cap is None:
                    print("重新连接失败，退出...")
                    break
                consecutive_error_count = 0
            continue

    # 确保释放资源
    cap.release()
    cv2.destroyAllWindows()
    print(f"\n处理完成！共处理 {processed_count} 帧，保存 {saved_count} 张图片")


def get_color_for_class(cls_id):
    """为不同类别分配不同颜色"""
    color_map = {
        2: (0, 255, 0),  # car - 绿色
        5: (255, 0, 0),  # bus - 蓝色
        7: (180, 105, 255)  # truck - 标准粉红
    }
    return color_map.get(cls_id, (255, 255, 255))


def save_segmented_image(segmented_img, frame_count, output_dir, prefix="combined"):
    """保存带有车辆和道路标记的图片"""
    if segmented_img is not None:
        # 生成带时间戳的文件名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
        filename = f"{prefix}_frame_{frame_count}_{timestamp}.jpg"
        filepath = os.path.join(output_dir, filename)

        # 保存图片（使用高质量设置）
        cv2.imwrite(filepath, segmented_img, [cv2.IMWRITE_JPEG_QUALITY, 95])
        print(f"保存图片: {filename}")


def display_results(segmented_img, frame_count, vehicle_count, violations):
    """实时显示检测结果"""
    if segmented_img is not None:
        # 添加帧计数信息
        display_frame = segmented_img.copy()

        # 根据是否检测到车辆显示不同信息
        if vehicle_count > 0:
            status_text = f'Frame: {frame_count} - Vehicles: {vehicle_count}'
            status_color = (0, 255, 0)  # 绿色
        else:
            status_text = f'Frame: {frame_count} - No Vehicles'
            status_color = (0, 0, 255)  # 红色

        cv2.putText(display_frame, status_text, (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, status_color, 2)

        # 显示压线告警信息
        if violations:
            violation_text = f'VIOLATIONS: {len(violations)} vehicles detected crossing lines!'
            cv2.putText(display_frame, violation_text, (10, 60),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)

        # 显示图例
        legend_y = 100
        cv2.putText(display_frame, 'Legend:', (10, legend_y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
        legend_y += 30

        # 车辆类别图例
        for cls_id, cls_name in CLASS_NAMES.items():
            color = get_color_for_class(cls_id)
            color_bgr = (int(color[0]), int(color[1]), int(color[2]))
            cv2.putText(display_frame, f'{cls_name}: {color_bgr}', (10, legend_y),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, color_bgr, 2)
            legend_y += 25

        # 道路标记图例
        cv2.putText(display_frame, 'Lane Markings: Red', (10, legend_y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
        legend_y += 25

        # 压线图例
        cv2.putText(display_frame, 'Violating Vehicle: Red (thick border)', (10, legend_y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
        legend_y += 25

        # ROI区域图例
        cv2.putText(display_frame, 'ROI Region: Cyan border', (10, legend_y),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2)
        legend_y += 25

        # 检测区域图例
        if SHOW_DETECTION_BOX:
            cv2.putText(display_frame, 'Detection Region: Cyan box', (10, legend_y),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2)

        # 显示结果 - 现在会显示在1280x600的窗口中
        cv2.imshow('Vehicle + Road Markings Detection', display_frame)


def process_single_image(image_path, output_dir):
    """处理单张图片"""
    print("开始预测单张图片...")
    print(f"输入图片: {image_path}")
    print(f"输出目录: {output_dir}")

    # 读取图片
    frame = cv2.imread(image_path)
    if frame is None:
        print(f"错误：无法读取图片 {image_path}")
        return

    h, w = frame.shape[:2]
    # 默认梯形ROI坐标 (可根据需要调整)
    top_left = (w // 4, h // 2)  # 左上角
    top_right = (3 * w // 4, h // 2)  # 右上角
    bottom_right = (w, h)  # 右下角
    bottom_left = (0, h)  # 左下角

    #top_left = (1300, 345)  # 左上角
    #top_right = (1450, 345)  # 右上角
    # bottom_right = (1881, 732)  # 右下角
    # bottom_left = (483, 732)  # 左下角
    #bottom_right = (w, h)  # 右下角
    #bottom_left = (0, h)  # 左下角


    # 允许用户自定义梯形ROI坐标
    print("请输入梯形ROI的4个顶点坐标 (格式: x,y):")
    print(f"当前图片尺寸: {w}x{h}")
    print("提示: 坐标格式为 x,y (例如: 100,200)")

    try:
        top_left_input = input(f"左上角坐标 (当前: {top_left}) - 直接回车使用默认值: ").strip()
        if top_left_input:
            x, y = map(int, top_left_input.split(','))
            top_left = (x, y)

        top_right_input = input(f"右上角坐标 (当前: {top_right}) - 直接回车使用默认值: ").strip()
        if top_right_input:
            x, y = map(int, top_right_input.split(','))
            top_right = (x, y)

        bottom_right_input = input(f"右下角坐标 (当前: {bottom_right}) - 直接回车使用默认值: ").strip()
        if bottom_right_input:
            x, y = map(int, bottom_right_input.split(','))
            bottom_right = (x, y)

        bottom_left_input = input(f"左下角坐标 (当前: {bottom_left}) - 直接回车使用默认值: ").strip()
        if bottom_left_input:
            x, y = map(int, bottom_left_input.split(','))
            bottom_left = (x, y)

    except ValueError:
        print("输入格式错误，使用默认坐标")

    print(f"使用梯形ROI坐标: 左上{top_left}, 右上{top_right}, 右下{bottom_right}, 左下{bottom_left}")

    # 初始化ROI掩码
    global roi_mask
    roi_mask = create_trapezoid_roi_mask(frame.shape, top_left, top_right, bottom_right, bottom_left)

    # 获取道路标记掩码
    lane_marking_mask = get_road_segmentation_mask(frame)

    # 使用YOLOv8-seg进行车辆检测
    results = model.predict(
        source=frame,
        imgsz=1280,
        conf=0.6,
        iou=0.5,
        retina_masks=True,
        save=False,
        verbose=False
    )

    # 处理检测结果并合并道路标记掩码
    visualized_frame, vehicle_count, violations = combine_vehicle_and_road_masks(frame, results[0], lane_marking_mask)

    # 检查压线情况
    if violations:
        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        print(f"[{current_time}] 检测到压线告警:")
        for violation in violations:
            print(f"  - {violation['class_name']} 压线! 位置: ({violation['bbox'][0]}, {violation['bbox'][1]})")
            print(f"    检测区域: {violation['detection_bbox']}")
            # 如果有压线情况，保存图片

    # 保存结果
    if vehicle_count > 0:
        # 生成输出文件名
        original_name = os.path.splitext(os.path.basename(image_path))[0]
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_filename = f"combined_{original_name}_{timestamp}.jpg"
        output_path = os.path.join(output_dir, output_filename)

        # 保存合并结果
        cv2.imwrite(output_path, visualized_frame, [cv2.IMWRITE_JPEG_QUALITY, 95])
        print(f"检测到 {vehicle_count} 辆车辆，图片已保存: {output_filename}")

        # 创建显示窗口并设置大小
        cv2.namedWindow('Vehicle + Road Markings Detection', cv2.WINDOW_NORMAL)
        cv2.resizeWindow('Vehicle + Road Markings Detection', 1280, 600)

        # 显示结果
        display_results(visualized_frame, 0, vehicle_count, violations)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
    else:
        print("未检测到车辆，不保存图片")


def main():
    """主函数"""
    global should_exit

    # 注册信号处理器
    signal.signal(signal.SIGINT, signal_handler)

    # 创建带时间戳的输出目录
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir = f'results/{timestamp}'
    os.makedirs(output_dir, exist_ok=True)

    print("=" * 60)
    print("车辆 + 道路标记联合检测系统")
    print("=" * 60)
    print(f"车辆目标类别: {[CLASS_NAMES[cls_id] for cls_id in VEHICLE_CLASSES]}")
    print("道路标记类别: 类别3 (Lane Markings)")
    print("车辆用彩色掩码显示，道路标记用红色显示")
    print("压线车辆用红色显示并加粗边框")
    print("在车辆底部创建矩形检测区域检测压线")
    print(f"显示检测区域矩形框: {SHOW_DETECTION_BOX}")
    print("添加自定义梯形ROI检测区域")
    print("=" * 60)

    print("选择检测模式:")
    print("1. RTSP视频流检测")
    print("2. 单张图片检测")

    choice = "2"  # input("请输入选择 (1 或 2): ").strip()

    if choice == "1":
        # RTSP视频流检测
        rtsp_url = RTSP_URL
        process_interval = 10  # input("请输入处理间隔帧数 (默认10): ").strip()

        print(f"\n开始RTSP视频流检测...")
        print(f"RTSP地址: {rtsp_url}")
        print(f"处理间隔: 每{process_interval}帧处理一次")
        print(f"推理尺寸: 1280 (高精度)")
        print(f"输出目录: {output_dir}")
        print(f"压线阈值: {VIOLATION_THRESHOLD * 100}%")
        print(f"压线告警冷却时间: {VIOLATION_COOLDOWN}秒")
        print(f"车辆底部区域高度比例: {BOTTOM_REGION_HEIGHT_RATIO * 100}%")
        print(f"车辆底部区域宽度比例: {BOTTOM_REGION_WIDTH_RATIO * 100}%")
        print(f"显示检测区域矩形框: {SHOW_DETECTION_BOX}")
        print("按 'q' 退出检测 或 关闭窗口退出")
        print("-" * 40)

        process_rtsp_stream(rtsp_url, output_dir, process_interval)

    elif choice == "2":
        # 单张图片检测
        image_path = input("请输入图片路径 (直接回车使用默认路径): ").strip()
        if not image_path:
            image_path = r'H:\xiaomi\test\192.168.1.64_01_20251027093805446_frame_000087.png'

        process_single_image(image_path, output_dir)

    else:
        print("无效选择！")

    # 确保清理资源
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()